Dynamically Pruning SegFormer for Efficient Semantic Segmentation

SegFormer aims to extract hierarchical feature representations at multiple scales for semantic segmentation. SegFormer has a Mix Transformer (MiT) encoder followed by a dense MLP decoder, as shown in the left side of Figure 1. The encoder of SegFormer (i.e., MiT) has four different stages. At stage-$i$ of MiT, the input feature map $\bar{\mathbf{X}}_{i}\in\mathbb{R}^{C_{i}\times H\times W}$ is transformed to a patch embedding $\mathbf{X}_{i}=\textrm{PatchMerge}(\bar{\mathbf{X}}_{i})\in\mathbb{R}^{N\times C_{i}}$, where $N$ is the sequence length, and $C_{i}$ is the hidden dimension for stage-$i$. The efficient multi-head attention (MHA) of each transformer layer then computes the query $\mathbf{Q}=\textrm{FC}(\mathbf{X}_{i})\in\mathbb{R}^{N\times C_{i}}$, key $\mathbf{K}=\textrm{FC}(\mathbf{X}_{i})\in\mathbb{R}^{\bar{N}\times C_{i}}$ and value $\mathbf{V}=\textrm{FC}(\mathbf{X}_{i})\in\mathbb{R}^{\bar{N}\times C_{i}}$, where $\textrm{FC}(\cdot)$ is the linear layer. Note that both $\mathbf{K}$ and $\mathbf{V}$ have smaller sequence lengths $\bar{N}=N/R$ so as to reduce the time complexity of self-attention to $O(N^{2}/R)$. The output of MHA $\mathbf{X}_{mha}$ thus can be obtained by

	$\displaystyle\mathbf{A}=\textrm{softmax}(\frac{\mathbf{Q}\mathbf{K}^{\top}}{\sqrt{\bar{N}}})\mathbf{V}\in\mathbb{R}^{N\times C_{i}},$		(1)
	$\displaystyle\mathbf{X}_{mha}=\textrm{LN}\big{(}\hat{\mathbf{X}}+\textrm{FC}(\mathbf{A})\big{)}\in\mathbb{R}^{N\times C_{i}},\vspace{-0.5ex}$		(2)

where $\textrm{LN}(\cdot)$ denotes the layer-normalization layer. Then the mix feed-forward network (Mix-FFN) transforms $\mathbf{X}_{mha}$ by:

$$\mathbf{X}_{ffn}=\textrm{LN}(\textrm{FC}(\textrm{GeLU}(\textrm{Conv}(\textrm{FC}(\mathbf{X}_{mha})))+\mathbf{X}_{mha}),\vspace{-0.5ex}$$

(3)

where $\textrm{Conv}(\cdot)$ denotes the convolutional operations. Notably, SegFormer does not employ positional encodings inside the image patches but uses a convolutional layer $\textrm{Conv}(\cdot)$ to leak location information[16].

Finally, the MLP decoder transforms and concatenates the stage-wise feature maps $\{\mathbf{X}_{i}\}_{i=1}^{4}$ together to make the pixel-wise predictions. More details can be found in[6].

The hierarchical visual representations in SegFormer also lead to the increasing width in the MiT encoder, which brings more computational burden on low-power devices.

However, we find that each neuron is not always informative across different input instances. Specifically, we take a trained SegFormer with MiT-B0 encoder to conduct inference over 2,000 instances of the ADE20K benchmark and observe neuron magnitudes across all testing instances. We show the box plot of neuron magnitudes of a certain MiT layer on the right side of Figure 1 The box plot contains the median, quartiles, and min/max values, which can help categorize these neurons into three types: Type-I: large median with a small range, meaning that the neuron is generally informative across different instances; Type-II: large median with a large range, meaning that the information of neuron highly depends on the input; Type-III: small median with a small range, meaning that these neurons are mostly non-informative regardless of different input.

Given the above observation, it would be promising to reduce the computation inside the MiT encoder by identifying and pruning Type-II and Type-III neurons based on the input.

In order to identify and prune the un-informative neurons based on the input, we propose Dynamic Gated Linear (DGL) layer, a plug-and-play module to substitute the original linear layer in SegFormer. The DGL structure is shown in Figure 2.

The workflow of the DGL layer is as follows. Given input $\mathbf{X}\in\mathbb{R}^{N\times C}$, the DGL layer computes the gate $\mathbf{M}=g_{\phi}(\mathbf{X})\in\mathbb{R}^{N\times\bar{C}}$ via a gate predictor $g_{\phi}(\cdot)$ parameterized by $\phi$, where is $\bar{C}$ the output dimension. The gate $\mathbf{M}$ can be then applied to mask both the output dimensions of the current linear layer parameter $\mathbf{W}_{1}\in\mathbb{R}^{C\times\bar{C}}$, as well as the input dimensions of the next layer parameter $\mathbf{W}_{2}\in\mathbb{R}^{\bar{C}\times C}$. Therefore, the computation can be reduced for two consecutive layers as:

$$\mathbf{Y}=\mathbf{X}(\mathbf{W}_{1}\circ\mathbf{M})\in\mathbb{R}^{N\times\bar{C}},\ \ \mathbf{Z}=\mathbf{Y}(\mathbf{W}_{2}\circ\mathbf{M}^{\top})\in\mathbb{R}^{N\times C},$$

where $\circ$ denotes the element-wise product.

The design of gate predictor $g_{\phi}(\cdot)$ is the key to achieve dynamic pruning. As the input the gate predictor, we first summarize the sequence by average pooling over the $N$ image patches of input $\mathbf{X}$, followed by layer-normalization to scale it back to normal ranges, i.e., $\hat{\mathbf{X}}=\textrm{LN}(\textrm{AvgPool}(\mathbf{X}))\in\mathbb{R}^{C}$. Intuitively, aside from the input $\mathbf{X}$, the parameter $\mathbf{W}$ should be also incorporated to determine which neuron to prune in the output. We thus feed both $\hat{\mathbf{X}}$ and $\mathbf{W}_{1}$ into a two-layer $\textrm{MLP}(\cdot)$ activated by $\textrm{ReLU}(\cdot)$ to obtain the mask $\mathbf{M}$ as

$\displaystyle\mathbf{M}=\textrm{Top-r}(\mathbf{G}),\quad\mathbf{G}=\textrm{MLP}(\hat{\mathbf{X}}\mathbf{W}_{1})\in\mathbb{R}^{\bar{C}},\vspace{-0.5ex}$

where $\textrm{Top-r}(\cdot)$ keeps the top $r\%$ percentage largest elements in MLP output logits $\mathbf{G}$, and zero out the rest ones. In order to achieve a smooth transition to sparse parameters, we adopt an annealing strategy by gradually increasing the sparsity ratio at the $t$-th step as $r_{t}=r\min(1,\frac{t}{T})$, where $T$ is total annealing steps. Note that the $\textrm{Top-r}(\cdot)$ operation inevitably introduces information loss. To remedy this, we also encourage sparsity over the $\mathbf{G}$, which cab be achieved by $\ell_{1}$ norm penalty over $\mathbf{G}$ as the following:

$$\vspace{-0.5ex}\mathcal{L}_{m}=\lambda_{m}\sum\|\mathbf{G}\|_{1}.\vspace{-0.5ex}$$

(4)

We deploy the DGL layer to prune $\mathbf{Q}$, $\mathbf{K}$ and $\mathbf{V}$, as well as the intermediate layer of the Mix-FFN of the SegFormer, as a majority of computation lies in the MiT Encoder. For the MLP decoder, we replace the concatenation by addition over $\{\mathbf{X}_{i}\}_{i=1}^{4}$, so that the computation can be further reduced.

Remark The gate predictor $g_{\phi}(\cdot)$ has only $O(C\bar{C})$ computational complexity, which is negligible compared with the $O(NC\bar{C})$ complexity in the linear layer. However, all model parameters together with gate predictors should be saved, as every single parameter can be potentially activated depending on the input. Finally, dynamic pruning has also been previously explored in[17] for CNNs, where they only pass input $\mathbf{X}$ to the gate predictor, missing the information inside the parameter $\mathbf{W}_{1}$ for the gate prediction.

The dynamically pruned SegFormer inevitably loses the knowledge from the original model. To bridge the gap incurred by dynamic pruning, we adopt two-stage knowledge distillation[15, 18], which demonstrates superior performance on transformer-based models. In the first stage, we minimize the mean square error (MSE) between the student output $\tilde{\mathbf{X}}_{i}$ and the teacher output $\mathbf{X}_{i}$ at each SegFormer block, i.e.,

$$\ell_{1}=\sum_{i=1}^{4}\textrm{MSE}(\tilde{\mathbf{X}}_{i},\mathbf{X}_{i}).\vspace{-1ex}$$

(5)

Afterwards, with the student logits $\tilde{\mathbf{Y}}$, we minimize the conventional cross entropy loss (CE) with ground truth data $\mathbf{Y}^{\star}$ and soft-cross entropy (SCE) with teacher logits $\mathbf{Y}$ as:

$$\ell_{2}=\textrm{CE}(\tilde{\mathbf{Y}},\mathbf{Y}^{\star})+\lambda_{s}\textrm{SCE}(\tilde{\mathbf{Y}},\mathbf{Y}).\vspace{-1ex}$$

(6)

Note that for both two stages, we incorporate the sparsity regularization in Equation (4).

We empirically verify the proposed approach on ADE20K[19] and Cityscapes[20], two benchmark dataset for semantic segmentation. We follow the standard data pre-processing in[6], i.e., the images of ADE20K and Cityscapes are cropped to $512\times 512$ and $1024\times 1024$ respectively. We take the batch size as $16$ on ADE20K, and $8$ on Cityscapes, both of which are trained over 8 NVIDIA V100 GPUs. As we compare both the performance and efficiency of the model, we report the mean Intersection over Union (mIoU) together with computational FLOPs (G).

For the main experiments in Section 4.2, we present results for both real-time settings and non real-time settings. We adopt MiT-B0 as the encoder backbone for the real-time setting, and MiT-B2 for the non real-time setting. The dynamic pruning is based on the released model checkpoints ¹¹1https://github.com/NVlabs/SegFormer.. During the training, we take the sparsity regularizer $\lambda_{m}=0.005$, and soft cross-entropy regularizer $\lambda_{s}=0.5$ throughout the experiments. The training iterations are set to $160$K, where the first $50\%$ steps are used for sparsity annealing. We keep the rest configurations consistent with ones used in SegFormer, e.g., with learning rate of $0.00006$ under the “poly” LR schedule, and without auxiliary losses and class balance losses. We name the dynamically pruned SegFormer as DynaSegFormer.

We present the main results in Table 1, where $30\%$ and $50\%$ neurons are dynamically pruned in the MiT encoder. It can be found that for both the real-time and non real-time settings on the two benchmarks, our DynaSegFormer can easily outperform previous CNN-based models with significantly fewer computational FLOPs. For instance, our DynaSegFormer outperforms DeepLabV3+ by $2.9\%$ mIoU with only $5\%$ of its FLOPs. While the DGL layer increases the size of SegFormer by around $12\%$, it allows to identify and prune the instance-wise redundant neurons given different input. For example, we can achieve $36.9\%$ mIoU on ADE20K for the real-time setting, which is only $0.5\%$ inferior to SegFormer using only $40\%$ of the original computational FLOPs.

For ablation studies, we adopt a MiT-B0 as the encoder under 50% sparsity for pruning. The model is trained with $40$K iterations for fast verification, and results are present in Table 2. We analyze the following aspects: 1) Knowledge Distillation: Comparing with training without distillation (row #4), two-stage distillation (row #1) can improve the mIoU by around $1.2\%$. However, distillation with only stage-1 or stage-2 cannot bring such improvement (rows #2, #3). The finding is consistent with[15, 18] that the technique can better transfer the learned knowledge to the compressed model. 2) Annealing Sparsity: It is found generally helpful to anneal the sparsity during the training. For instance, it can improve the mIoU by $0.7\%$ for training with distillation (rows #1, #5), and $0.5\%$ for training without distillation (row #4, #6). The observations match our intuition since a smooth transition of sparsity ratio can better calibrate the parameters throughout the training process. 3) Dynamic v.s. Static Pruning: Finally, we verify the advantages of dynamic pruning over static pruning. For static pruning, we follow the widely used magnitude-based pruning[24], where the sparse mask is invariant to different input. It can be found that even with both distillation and sparsity annealing, static pruning is still inferior to dynamic pruning by $2.9\%$ (row #1, #7), while the gap is even enlarged to $3.2\%$ when trained without distillation (row #4, #8).

In the next, we analyze how the learned DGL layer keeps the Type-I neurons, but identifies and prunes the Type-II and Type-III neurons, as mentioned in Section 2.2. We count the number of active gate logits generated by the DGL layer, i.e., $\sum_{i}\mathbb{I}(G_{i}>0)$ during the inference, where $\mathbb{I}(\cdot)$ is the indicator function. From Figure 3, it can be found that there are around $15$ neurons that are always activated in the $2,000$ testing instances on ADE20K, indicating that they are the Type-I neurons. The Type-II neurons can be selectively activated ranging from $10\sim 1988$ counts, while the Type-III neurons are not visible since they are always inactive throughout all instances. Consequently, by identifying the types of neurons, the dynamic scheme selectively prunes neurons to better explain the data under limited computational constraints.