module, which is the main cost of computations and latency (over 60%). To achieve this, we first formulate the problem with pruning parameterization to build a pruning framework with a soft mask as a representation of the pruning policy. With this soft mask, we further adopt thresholding to convert the soft mask into a binary mask so that the model is trained with actual pruned weights to obtain pruning results directly. This is significantly different from other methods [27,47] to train with unpruned small non-zero weights and use fine-tuning to mitigate the performance degradation after applying pruning. Besides, to update the soft mask as long as the pruning policy, we adopt to straight though estimator (STE) method to make the soft mask differentiable. Thus, we can build the pruning parameterization framework with minimal overhead.

Based on this framework, we need to search the bestsuited layer width for each layer in the token pyramid module. It is non-trivial to perform the search. As the token pyramid module needs to extract multi-scale information from multiple spatial resolutions, the large hierarchical search space leads to difficulties of convergence. To resolve this problem, we adopt a bi-level optimization method. In the outer optimization, we try to obtain the pruning policy based on the pruning parameters (the soft mask). In the inner optimization, the optimized model weights with the best segmentation performance under this soft mask can be obtained. Compared with a typical pruning method, our work incorporates the implicit gradients with second-order derivatives to further guide the update of the soft mask and achieve better performance. Our experimental results demonstrate that we can achieve 38.9 mIoU (mean classwise intersection-over-union) on the ADE20K dataset with a speed of 56.5 FPS on Samsung S21, which is the highest mIoU under the same computation constraint with real-time inference speed. As demonstrated in Figure 1, our models can achieve a better tradeoff between the mIoU and speed.

We summarize our contributions below, • We propose a pruning parameterization method to build a pruning framework with a soft mask. We further use a threshold-based method to convert the soft mask into the binary mask to perform actual pruning during model training and inference. Besides, STE is adopted to update the soft mask efficiently through gradient descent optimizers. • To solve the pruning problem formulated with the framework of pruning parameterization, we propose a bi-level optimization method to utilize implicit gradients for better results. We show that the second-order derivatives in the implicit gradients can be efficiently obtained through first-order derivatives, saving computations and memory. • Our experimental results demonstrate that we can achieve the highest mIoU under the same computation constraint on various datasets. Specifically, we can achieve 38.9 mIoU on the ADE20K dataset with a real-time inference speed of 56.5 FPS on the Samsung S21.

Real-Time Semantic Segmentation. Though semantic segmentation based on CNNs [6,22,33,58,71] can achieve great performance, it typically costs large amounts of computations with slow inference speed. Furthermore, with the ever-increasing popularity of edge devices such as mobile phones, it is necessary to achieve fast inference speed for semantic segmentation on edge devices [4]. Besides human designed lightweight models [21,40,67], neural architecture search (NAS) methods [1,7,42,45,69] are also adopted to search lightweight models. As a pioneer in handcrafted real-time segmentation, ENet [50] designs a lightweight model for fast inference. DeepLabV3+ [6] adopts atrous separable convolution to reduce computation counts and uses the lightweight Mo-bileNetV2 [56] as the backbone. BiSeNet [66,67] and STDC [21] utilizes a two-branch architecture, where the deep branch extracts spatial information, and the shallow branch learns details. SFNet [40] uses flow alignment module to fuse context and spatial information.

Inspired by the recent success of NAS, some works automatically search lightweight segmentation models. To reduce the computation cost, FasterSeg [7] incorporates latency regularization during search. DCNAS [69] uses a densely connected search space and employs a gradientbased direct search method. NASViT [24] proposes a gradient projection algorithm to deal with the gradient conflict issues and improve the convergence performance. HR-NAS [15] keeps high-resolution representations in its encoder and the search process to maintain high accuracy in dense prediction. RTSeg [43] redesigns backbone and proposes a latency-driven search framework.

Although many lightweight segmentation models are developed, it is still hard for them to run real-time inference on resource-and power-limited GPUs of edge devices. Vision Transformers. By utilizing the self-attention mechanism, ViTs [17] can achieve competitive results against CNN models in vision tasks. Inspired by the great success of ViTs, many research efforts are devoted to dense prediction tasks with ViTs on complex datasets. SETR [72] utilizes a ViT-based encoder to extract highlevel semantic information. Instead of per-pixel prediction, MaskFormer [11] performs mask-based prediction with a customize backbone and a transformer-based decoder. With a transformer decoder to explore masked attention, mask2Former [10] proposes a universal architecture and achieves SOTA performance for various segmentation tasks. MobileViT [48] and MobileFormer [9] mix CNN and ViT in their architectures, but they are not fast enough for real-time inference on edge devices. TopFormer [68] uti-lizes c CNN-based token pyramid module to extract multilevel tokens and lightweight ViT blocks to extract semantic information. It achieves low latency and high accuracy on complex datasets. Neural Architecture Search and Pruning. NAS and pruning are commonly used to find lightweight models and reduce the computation overhead. Given a search space, NAS tries to identify a superior model automatically. Reinforcement learning (RL) based NAS [51,75] and evolution-based NAS [18,53] usually need to train and evaluate each candidate model, leading to tremendous searching cost. To mitigate this, gradient-based NAS [5,12,29,46] is proposed to formulate a supernet with all candidate architectures and search the outstanding architecture through gradient descent methods. But it costs huge memory to incorporate all candidate architectures.

Network pruning is a compression technique to effectively reduce the DNN storage and computation cost [31]. In this work, we focus on structured pruning [39,61] to remove entire filters or channels of CONV layers, which can be accelerated effectively for fast inference.

In our method, we first formulate the pruning problem with pruning parameterization. Previous pruning methods usually depend on the magnitudes of model weights and adopt the in-differentiable sorting operations, leading to inconsistent performance after applying pruning with sorting and additional overhead with fine-tuning. To mitigate this, we propose a pruning parameterization method, which uses a soft mask (rather than magnitudes of weights) to indicate whether to prune and get rid of sorting operations. With the STE method [3], we are able to represent the pruning with parameters and directly train the pruning like a typical model training. Then we formulate our problem with pruning parameterization and introduce our solution.

We first introduce how to construct the soft mask. To accelerate the inference, we adopt channel pruning to search for a suitable width for each convolution (CONV) layer. Specifically, we insert a depth-wise CONV layer following each CONV layer that is supposed to be pruned as below,

where ⊙ denotes the convolution operation. w l ∈ R o×i×k×k is the weight parameters in l-th CONV layer, with o output channels, i input channels, and kernels of size k × k. a l ∈ R n×o×t×t ′ represents the output features of the l-th layer, with o channels and t × t ′ feature size. n denotes batch size. s l ∈ R o×1×1×1 is the weights of the depth-wise CONV layer. Each output channel of w l ⊙ a l-1 corresponds to one single element of s l . Thus s l can serve as a soft mask or pruning indicator for the l-th CONV layer to indicate whether to prune the corresponding channels.

Since s l is only a soft mask with continuous values, it can not represent the binary operation of pruning. To solve this, we adopt a threshold, and the forward pass with the mask is represented as

where b l ∈ {0, 1} o×1×1×1 is the binarized s l , and τ is a threshold which is simply set to 0.5 in our case. Each element of b l corresponds to one output channel of w l ⊙ a l-1 .

In the forward pass, we first convert the soft mask into a binary mask through a threshold and then perform the depth-wise CONV to perform actual pruning. Thus the output channels corresponding to the zero elements in b l are pruned, and the rest channels are preserved. We show the proof in Appendix A.

Advantages of the Binary Operation. With the binary operation to obtain the binary mask, the pruned weights become zero during computations of both training and inference. This is different from some pruning works [27,30,47] to perform training with small but non-zero weights and inference with binary weights, which has inconsistent performance between training and inference, and requires additional finetuning.

Why τ = 0.5. Since the s l parameters in Equation ( 2) are initialized randomly between 0 and 1, we set τ = 0.5 as a threshold for the binary operation. τ can be set to other values (such as 0.3 or 0.6). The key point is that s l can be updated to increase above τ to keep the corresponding channel or decrease below τ to prune the channel. The binary operation in Equation ( 2) is non-differential, leading to difficulties for back-propagation. To solve this, we propose to use STE [3,64] for back-propagation below,

where we directly pass the gradients of b l to s l so that we can update the soft mask. Why Named Pruning Parameterization. With the binarization and STE method, the pruning process can be represented with the soft mask s = {s l }. We can update the soft mask to update the pruning policy based on its gradients. Thus s is the pruning parameters to denote and control pruning, i.e., pruning parameterization.

Difference with Other Pruning Works. Based on pruning parameterization, we decouple the pruning policy from model parameter magnitudes so that pruning does not further depend on the weight magnitudes. Unlike previous works on pruning to force pruned weights to be or as close as to zeros [27,30,47], our method does not have such a constraint that pruned weights should be zero. Instead, once the corresponding binary mask is turned from 1 to 0, the information in pruned channels is preserved rather than zeroed out since zeros in b l can block gradient flow to the corresponding weights. As a result, pruned channels are free to recover and contribute to accuracy if their corresponding elements in b l switch from 0 to 1.

Difference with Other Mask Methods. Although some other works also adopt indicator/mask-based pruning such as [27,28,36,37], our method is more straightforward and effective. For example, unlike our method to train the soft mask with STE directly, DAIS [27] relaxes the binarized channel indicators to be continuous. To bridge the nonnegligible discrepancy between the continuous model and the target binarized model, it further uses an annealingbased procedure to steer the indicator convergence toward binarized states. Some works [36,65] adopts batchnorm (BN) layers and uses the cumulative density function (CDF) of a Gaussian distribution as the mask variable, with a CDFbased loss function and the Gumbel-Softmax trick to update the mask at the cost of additional random sampling and complex gradient revision. Besides, the works [23,65] create a mask to multiple the channels weights. This is different from our design with the depth-wise CONV operation for easy mask creation and direct mask training.

Based on the soft mask, we can train and prune the model with the following loss function,

where L(w, s) is the cross-entropy loss, and L reg is the regularization term related to the sparsity or pruning ratio. For simplicity, we take Multiply-Accumulate operations (MACs) as the regularization rather than parameter number to estimate the on-device execution cost more precisely. β can weight the loss and stabilize training. L reg can be defined as the squared ℓ 2 norm of the difference between current MACs and target MACs C,

where o ′ l is the number of remaining channels after pruning, and t l × t ′ l is the output feature size.

With pruning parameterization, the per-layer width search problem can be formulated as follows

It is a bi-level optimization problem [25,35]. In the inner optimization of Equation ( 7), we optimize the model parameters under a given soft mask with a commonly used squared ℓ 2 norm as a regularization to mitigate the overfitting problem. In the outer optimization of Equation ( 6), we optimize the soft mask to minimize the loss. Each time before updating s, we first need to obtain w * .

The objective is to find the soft mask (search a suitable width) for each layer to minimize the loss with optimized model weights. We adopt a bi-level pruning method to solve this problem. Compared with a typical gradient descent method to update the parameters with first-order derivatives, the bi-level optimization method incorporates implicit gradients with second-order derivatives to adjust the first-order term, leading to higher training efficiency with better convergence results. For easy of expression, since each channel has a corresponding mask, in this section, we broadcast the masks s and b so that the masks have the same dimension as the model weights w.

From the inner optimization, w * is a function of s and different s can lead to different w * . Thus, to minimize L m (w * , s) in Problem (6), we need to compute the gradients with reference to s as below, dLm(w * , s) ds = dw * ds ∇wLm(w * , s) + ∇sLm(w * , s), (8) where ∇ w and ∇ s denote the partial derivatives of the loss function with reference to w and s, respectively. dw * ds represents the vector-wise full derivative, and we omit the transpose expression. Since w * is implicitly defined as an optimization problem in Equation ( 7), dw * ds is also known as implicit gradients [52,54,55,70]. With g(w, s) = L(w, s) + 1 2λ w T w, dw * ds can be obtained through the following,

where ∇ 2 sw and ∇ 2 w are the second-order partial derivatives. We show the proof in Appendix B.

The Hessian matrix ∇ 2 w g(w * , s) can be given by

where we adopt a Hessian-free approximation that ∇ 2 w L(w * , s) = 0 as DNNs usually have piece-wise linear decision boundary with ReLU functions. Thus, Equation (8) can be transformed to

Compared with a typical gradient descent method, the bilevel optimization incorporates the second-order derivatives ∇ 2 sw L(w * , s)∇ w L m (w * , s) to adjust the first-order term ∇ s L m (w * , s).

It is difficult to obtain the second-order derivatives ∇ 2 sw L(w * , s). Usually certain approximation methods such as finite difference [8,20] may be adopted to save computation cost. But here we show that in this specific problem, we can directly obtain the analytical solution with first-order derivatives, which greatly saves computation efforts without approximation. Note that each channel has its corresponding mask, and we denote the masked channels as w b = w * • b where • means the element-wise multiplication and b is defined in Equation ( 2). We can obtain that

where diag(•) represents formulating a diagonal matrix with the diagonal vector. We show the proof in Appendix C.

Then we can transform Equation ( 11) into the following,

Thus although the implicit gradients incorporate secondorder derivatives, it can be analytically expressed with firstorder derivatives, greatly saving computation cost without any approximations.

In Equation ( 13), we need to create two copies w b and w to obtain their derivatives, which are still not memory efficient. To further reduce the complexity, we can obtain ∇ w b L(w b ) in the following,

The pruned weights (b = 0) do not contribute to the loss, so their gradients are zero. The difference between L(•) and L m (•) is just the regularization term L reg (•), which only relates to the sparsity and does not care the weight values. So ∇ w b L(w b ) and ∇ w L m (w) are equal if their weights are not pruned (b = 1). Combining Equation ( 13) and ( 14), we can obtain the following,

We can see that there is no need to keep a copy and compute the gradients of w b , thus saving memory. During practical implementation, since s and b are the channel-wise masks, we accumulate the gradients of all weights in each channel to update the channel mask following Equation (15). The Case without Implicit Gradients. In Equation ( 8), we can omit the first term with implicit gradients. Thus we have dLm(w * ,s) ds = ∇ s L m (w * , s) and there is no need to deal with the second-order term. It is easier to solve. But we demonstrate in the experiments that our method with implicit gradients can help to boost the performance without a significant increase of the computation cost in Section 5.4.

In each iteration, our framework has two steps, including the model weights training step and the mask updating step. In the first step, we update the weights w with a few training steps for a fixed mask s. Next in the second step, we update s with implicit gradients following Equation (15). Then we move on to the next iteration. In each step, we only update w or s without changing the other parameter. We discuss the advantages of the proposed method in the following. No Need for Finetuning. After training with the proposed method, we can obtain the final layer-wise pruning policy following s and the sparse model. Since the model is already trained with the binary masks, it can achieve good segmentation performance without further fine-tuning. First-Order Optimization. During the computation, we only adopt first-order optimization. Though we incorporate second-order derivatives in implicit gradients, we show that it can be analytically expressed with first-order derivatives in Equation ( 12), which greatly saves computation cost. We further optimize the computation process in Equation ( 15) to save memory cost. Recoverable Contribution. During pruning, though some channels are pruned, their weights are not set to 0 due to the protection of the mask. When their mask is updated from 0 to 1, they can recover and contribute the accuracy.

ADE20K. ADE20K [73,74] is a scene parsing dataset containing 25k images in the training set and 2k images in the validation set with 150 label classes. Cityscapes. Cityscapes [14] is a dataset of urban street scenes from cars collected in 50 cities. It includes 5,000 finely annotated images, in which 2,975 images are used for training, 500 for validation, and 1,525 for testing. We exclude the extra training data with coarse labels throughout this paper. This dataset has 30 label classes, and 19 of them are used for segmentation. The resolution of images is 2048 × 1024. Cityscapes dataset is an intensively studied benchmark for semantic segmentation, but it is challenging to perform real-time inference on such a high resolution. Pascal VOC. PASCAL Visual Object Classes (VOC) 2012 [19] is a widely used dataset for semantic segmentation, classification, and object detection tasks. There are 1,464 images for training and 1,449 images for validation. We show the results on Pascal VOC in Table 4. Evaluation Metrics. For semantic segmentation evaluation, we use the mean of class-wise intersection-over-union (mIoU) to measure the accuracy performance. We introduce the details of mIoU in Appendix D. For our results, we run our algorithm three times and report the mean and variance of the mIoU performance. For the baseline methods, some Table 1. Comparison of our searched model and prior arts on the ADE20K val dataset. We compare with popular handcraft baselines in the first segment, NAS-based models in the second segment, pruning-based methods in the third segment and lightweight ViT-based models in the fourth segment. We measure the FPS on the Qualcomm Adreno 660 GPU of the Samsung Galaxy S21 mobile phone. Some FPS results are not available due to unsupported operations on the mobile device. methods cost too many resources and we can hardly rerun their experiments. Some methods provide their well-trained models, and we can test with the trained model.

Train Settings. We perform the pruning to search for a suitable width for each CONV layer in the TopFormer model.

To enable a larger search space, we adopt the TopFormer architecture and use a larger per-layer width compared with the TopFormer-Base model. So our unpruned model has 4.1GMACs (with 39.9 mIoU), larger than the 1.8GMACs of the TopFormer-Base model. We use stochastic gradient descent (SGD) optimizer, and momentum is set to 0.9, and set the batch size to 8 on each GPU. For ADE20K, the initial learning rate is set to 1.2 × 10 -4 , and the "poly" learning rate policy is applied. For the Cityscapes dataset, the initial learning rate is set to 3 × 10 -4 , and we apply the "poly" learning rate policy. For PASCAL VOC 2012, we set the initial learning rate as 0.01. Learning rate value is determined as 1 -iter total iter 0.9

where iter refers to the current iteration number. For the ADE20K dataset, we incorporate data augmentation by resizing with the random ratio between 0.5 and 2.0 as well as random flipping. On the Cityscapes dataset, multiple random scaling {0.5, 0.75, 1.0, 2.0} and fixed size cropping of 512 × 1024 are adopted for data augmentation. We choose the crop size for a better trade-off between mobile capacity and accuracy. We set hyperparameters β = 0.01 and λ = 0.1 in the experiments. Test Settings. Instead of muti-scale testing, we employ single scale testing for a fair comparison. For the ADE20K dataset, we use 512 × 512 as the input resolution. For the Cityscapes dataset, 512 × 1024 (rather than 1024 × 2048) is used as the inference resolution during testing for the following reasons. (i) In practice, we cannot use the resolution 1024 × 2048 since it causes memory overflow problems on our selected mobile phone. (ii) Besides, since the screens on edge devices such as mobile phones are not very large, the resolution of 512 × 1024 is good enough to serve on the small screens. (iii) Moreover, we find that this 512 × 1024 resolution can greatly speedup the inference on the mobile phones without significant accuracy degradation. Experiment Environments. We train and prune the model using PyTorch 1.9 and CUDA 11.1 on 8 NVIDIA RTX TI-TAN GPUs. We measure the mobile latency on the GPU of an Samsung Galaxy S21 smartphone, with Qualcomm Snapdragon 888 mobile platform integrated with Qualcomm Kryo 680 Octa-core CPU and a Qualcomm Adreno 660 GPU. Note that for most baseline works, even the reduced resolution (512 × 1024) can cause an out-of-memory problem on the selected mobile device. Compiler Framework on Mobile Devices. We need to compile the models before they can be executed on mobile devices. For TopFormer, we adopt the compiler TNN [13] to report the speed performance, which is also used in the original TopFormer paper. To further improve the inference speed, we adopt several compiler optimization methods and Table 2. Our search results on the Cityscapes val dataset. We compare with popular handcraft baselines, NAS-based models, pruning based methods and lightweight ViT-based models. We measure the FPS on the Qualcomm Adreno 660 GPU of the Samsung Galaxy S21 mobile phone. Some FPS results are not available due to unsupported operations on the mobile device. develop an advanced compiler framework to compile our models and test their on-mobile speed. We show the details about our compiler optimization in Appendix E.

Segmentation Performance. Based on the TopFormer model, we obtain three models (Ours-Base, Ours-Small, and Ours-Tiny) with different computations. We show the comparison results on ADE20K in Table 1 and Cityscapes in Table 2. (i) For ADE20K, we can observe that the humandesigned segmentation models and NAS-based models usually consume many computations (such as DeepLabV3+ with 26.9GMACs) in terms of MACs (multiply-accumulate operations). They can hardly achieve real-time inference on edge devices. Some NAS-based models with a small number of computations are not able to achieve high mIoU (such as HR-NAS-A with 33.2 mIoU). (ii) For the comparison with other pruning methods, we start from the same dense model and set the target GMACs to the same number (1.2GMACs) to make a fair comparison. We can see that our small model can achieve higher mIoU given the same GMACs. (iii) Compared with the ViT-based TopFormer, our models can achieve higher mIoUs with non-trivial improvements under the same computation budget (such as Ours-Small 37.5 mIoU v.s. 36.1 of TopFormer-Small under the same 1.2GMACs). Our models can achieve a faster inference speed (FPS higher than 50) on mobile phones compared with TopFormer models. We show the comparison with the baselines in terms of mIoU and FPS in Figure 1.

We can achieve a better trade-off between mIoU and FPS. (i) On the Cityscapes dataset, compared with handcrafted CNN-based segmentation models, including BiSeNetV2 and STDC, our searched base model greatly reduces the GMACs (such as Ours-Base 3.6GMACs v.s. 24.6GMACs of BiSeNetV2) and achieves non-trivial better accuracy (Ours-Base 74.7 mIoU v.s. 73.4 mIoU of BiSeNetV2), as shown in Table 2. (ii) Compared with NAS-based methods such as FasterSeg, our models can achieve higher mIoU with much smaller computation costs. Other NAS methods consume too many computations (e.g., DCNAS with 300GMACs) to be deployed on edge devices. (iii) Compared with other pruning baselines, under the same computation budgets (2.4GMACs), our small model can achieve higher mIoU. (iv) Compared with transformerbased methods such as SegFormer-B0, similarly, our models achieve higher mIoU with less computations. Our small and tiny models have similar computations compared with TopFormer-Base and TopFormer-Tiny. But we can achieve higher mIoU. Speed Performance on Mobile Devices. Our base model can achieve 56.5 FPS on the mobile device (Samsung Galaxy S21), which implements real-time execution with competitive segmentation performance, as shown in Table 1. Other baseline methods except TopFormer can hardly achieve real-time segmentation on edge devices, usually with FPS lower than 10. Our faster inference speed than TopFormer is achieved with the compiler optimization techniques, detailed in Appendix E. Search Overhead. We show the search cost in Table 3. We only show the cost of our small model since our base, small and tiny models have similar search costs. It takes approximately 1.3 GPU days, which is smaller than the search cost of most other NAS-based segmentation methods. Our method can efficiently search out a compact model with Our method is general and can be applied to other model structures. We show the results for DeepLabV3+ [6] on Cityscapes in Table 5. Our two searched models have fewer parameters and computations with accelerated on-mobile inference speed, while achieving better mIoU compared with the original DeepLabV3+ model. Other models such as BiSeNetV2 or STDC cost much more computations.

In our bi-level optimization, we incorporate implicit gradients in Equation (11). To demonstrate the advantages with implicit gradients, we consider the case without implicit gradients which omits the second term in Equation (11) with just dLm(w * ,s) ds = ∇ s L m (w * , s). We compare the performance of the solution without implicit gradients and our bilevel solution with implicit gradients on Cityscapes in Table 3. We can observe that our solution with implicit gradients has a search cost slightly higher than the solution without implicit gradients, demonstrating that our first-order solution for the second-order derivatives in Equation ( 12) can effectively save computation cost. For mIoU, our method can achieve higher mIoU with non-trivial improvements, demonstrating that incorporating implicit gradients can effectively boost the performance. Hyperparameter Tuning. We show the results with various β and λ in Appendix G. To compare with baselines under certain computations, we mainly show the results of our three sparse models (Ours-Base, Ours-Small, and Ours-Tiny). We show the results of more models under other computations in Appendix H.

We propose pruning parameterization with the thresholding and STE methods to build a pruning framework. Based on the framework, we formulate the problem and propose a bi-level optimization method with the implicit gradients. Our experimental results demonstrate that we can achieve the highest mIoU under the same computation constraint on various datasets. Specifically, we can achieve 38.9 mIoU on the ADE20K with a real-time inference speed of 56.5 FPS on the Samsung S21.