Layer Pruning on Demand with Intermediate CTC

Connectionist temporal classification (CTC) solves the sequence prediction problem by introducing a monotonic alignment. For the encoder output $x\in\mathbb{R}^{T\times D}$ of length $T$ and feature dimension $D$, the CTC layer computes the likelihood of target sequence $y$:

$$P(y|x):=\sum_{a\in\beta^{-1}(y)}P(a|x),$$

(1)

where $\beta^{-1}(y)$ is the set of possible alignments that are compatible to $y$ and are of length $T$, and $a$ is an alignment in the set. The probability of an alignment is modeled as a factorized distribution:

	$\displaystyle P(a|x)$	$\displaystyle:=\prod_{t}P(a[t]|x[t])$		(2)
	$\displaystyle P(a[t]|x[t])$	$\displaystyle:=\text{Softmax}(\text{Linear}(x[t]))_{a[t]}$		(3)

where $a[t]$ and $x[t]$ mean the $t$-th value of $a$ and $x$ respectively.

During inference, the most probable alignment is found by greedy decoding, allowing fast and non-autoregressive inference.

Transformer [18] is a multi-layer architecture with self-attention and residual connection [34]. Transformer layer uses self-attention to learn global information and residual connection to help training of a deep neural network.

We use the encoder part of the Transformer architecture. Let $L$ be the number of layers in the encoder. The $l$-th layer computes the representation $x_{l}$ for given input $x_{l-1}$ by:

	$\displaystyle x^{\text{MHA}}_{l}$	$\displaystyle=\text{SelfAttention}(x_{l-1})+x_{l-1}$		(4)
	$\displaystyle x_{l}$	$\displaystyle=\text{FeedForward}(x^{\text{MHA}}_{l})+x^{\text{MHA}}_{l},$		(5)

where $x_{0}$ indicates the input of the Transformer encoder.

The final representation $x_{L}$ is then fed to CTC layer to minimize the following loss:

$$\mathcal{L}_{\mathsf{CTC}}:=-\log P_{\mathsf{CTC}}(y|x_{L}).$$

(6)

Stochastic depth [27, 28] is a regularization method designed for deep residual networks [34]. During training, each layer is randomly skipped or not with a given probability $p$. For each iteration, $u$ is sampled from a Bernoulli distribution so that the probability of $u=1$ is $p$ and the probability of $u=0$ is $1-p$. The output is computed by modifying Eqs. (4) and (5) as:

	$\displaystyle x^{\text{MHA}}_{l}$	$\displaystyle=\frac{u}{p}\cdot\text{SelfAttention}(x_{l-1})+x_{l-1}$		(7)
	$\displaystyle x_{l}$	$\displaystyle=\frac{u}{p}\cdot\text{FeedForward}(x^{\text{MHA}}_{l})+x^{\text{MHA}}_{l}.$		(8)

If $u=0$, the residual parts are skipped (i.e., $x_{l}=x_{l-1}$).

LayerDrop [5] is an application of stochastic depth to layer pruning. After training the model with stochastic depth, removing some layers from the model gives new smaller sub-model which also has reasonable performance without any fine-tuning.

Intermediate CTC [22] is an auxiliary loss designed for CTC modeling. It regularizes the model using an additional CTC loss attached at the intermediate layer of the encoder.

Let $l_{1},\ldots,l_{K}$ be the $K$ positions ($K<L$) of the intermediate layers. The intermediate loss is defined as:

$$\mathcal{L}_{\mathsf{InterCTC}}:=\frac{1}{K}\sum_{k}-\log P_{\mathsf{CTC}}(y|x_{l_{k}}).$$

(9)

Then, the training objective is defined by combining Eqs. (6) and (9):

$$\mathcal{L}:=(1-w)\mathcal{L}_{\mathsf{CTC}}+w\mathcal{L}_{\mathsf{InterCTC}}$$

(10)

with a hyper-parameter $w$.

Intermediate CTC has very small overhead during training, because the intermediate representation $x_{l_{k}}$ is naturally obtained when the final output $x_{L}$ is computed. The only additional cost is to compute CTC loss for the given representation, which is much smaller than the cost of the encoder. [22] explores various choices for the intermediate layer positions, and concludes that it is sufficient to use one layer ($K=1$) in the middle ($l_{1}=\lfloor L/2\rfloor$) for the regularization purpose. We revisit the effect of variants at Section 3.

Note that CTC and intermediate CTC share the same linear projection layer of Eq. (3), as intermediate CTC is treated as a regular CTC loss except that all of the encoder layers after the intermediate layer are skipped. This property is used at Section 4.1.

We first investigate how the two regularization techniques, stochastic depth and intermediate CTC, affect layer similarity. We train a 24-layer Transformer CTC model with four different configurations: (a) baseline, (b) with intermediate CTC ($w=0.3$, at 12th layer), (c) with stochastic depth, and (d) with both. Figure 1 shows the similarity matrices of the four models.

From Figures 1 (a) and (b), we observe that intermediate CTC increases similarity of layers between the intermediate one (12th) and the final one (24th), at the cost of decreasing similarity of the first layers a bit. Especially, the intermediate layer is significantly similar to last layer and its nearby layers.

On the other hand, Figure 1 (c) shows that stochastic depth evenly increases the similarity of the last half of layers. Especially we see the neighboring layers have very high similarities. This can be explained by that the layer should behave similarly regardless of whether its previous layer has been skipped or not.

Finally, Figure 1 (d) shows the combination of two regularizations greatly increases similarity among all layers between the intermediate layer and the last layer. This gives us an insight that the both regularizations would help pruning the layers.

The position variant of intermediate CTC is first explored by [22], concluding that its impact to the accuracy is minimal. However, we find the variant gives big difference on layer similarity and layer pruning.

First, we put an additional intermediate layer at the 6th layer, in addition to the 12th one. As shown at Figure 2 (a), it increases the similarity of layers below the 12th layer, compared to Figure 1 (d). This suggests that multiple branch variant affects the performance of pruning.

We then investigate the effect of the weight for the intermediate loss. [22] used $w=0.3$ as it is sufficient to enhance the model performance, but it may not be an optimal value for pruning, especially if there are multiple intermediate layers. Thus, we give equal weight to the two intermediate layers and the final layer, i.e., $w=0.66\approx 2/3$. Figure 2 (b) shows that the increase of intermediate weight helps overall similarity.

From the observations, we use the newly found hyper-parameters for the main experiments at Section 5.

The intermediate CTC loss is originally proposed for regularization, and it is not used during inference. However, in the context of layer pruning, the intermediate layer suggests a natural strategy to induce a smaller sub-model from the original model: remove all of the layers after the intermediate layer.

Moreover, we note that the method can be even extended to any depth, even if the corresponding layer is not explicitly trained with the intermediate CTC loss. It is because the linear projection layer of Eq. (3) is independent of specific layer and is shared by the intermediate and last layers, enabling that it can be used by other layers. Also, with the improved regularization presented at Section 3.2, we expect layers between the intermediate layer and the last layer give reasonable performance.

For the layers used as intermediate branches during training, we find the performance of their sub-models is as same as the one of the individually trained models of same depth. For the other layers, the performance is only slightly worse than one of individual models. We show the experimental result at Section 5.

LayerDrop [5] presents pruning strategies for stochastic depth regularization. The authors tried two strategies, removing every other layer and exhaustive search, and concluded the two strategies do not give significant difference. They also found removing consecutive layers causes significant degradation.

However, we find removing the last half gives nearly optimal performance, due to intermediate CTC. Especially, the intermediate sub-model is significantly better than removing every other layer. This suggests careful pruning strategy would be beneficial.

In this end, we suggest an iterative search augmented with the intermediate model. We start with the original model, and iteratively decrease model depth by one.

For brevity, we denote sub-models as the set of number indicating the layer of the original model. For example, if the 2-layer sub-model uses the 2nd and 4th layers of the original model, we denote the model as $\{2,4\}$.

For a given sub-model of depth $k$, we find a new sub-model of depth $k-1$ by the following steps:

1. 1.

   For each layer of the current model, we induce a new sub-model by removing the layer from the model. For example, from a given model $\{2,3,4\}$, we induce three sub-models: $\{2,3\},\{2,4\},\{3,4\}$.

2. 2.

   We also induce an intermediate sub-model $\{1,\dots,k-1\}$ from the original model. In the current example, the intermediate sub-model is $\{1,2\}$.

3. 3.

   Evaluate the induced sub-models using the validation set and choose the sub-model with the best accuracy.

Step 2 ensures that intermediate sub-models are always included in search space. Without Step 2, we found the iterative search may not discover them and sometimes perform worse.

Note that when a model is evaluated, its intermediate sub-models can be evaluated as well with very small additional cost. For example, evaluation of $\{2,3,4\}$ can be done along with evaluations of $\{2,3\}$ and $\{2\}$. This reduces computational cost for the search.

For layer pruning, we train a 24-layer Transformer model with intermediate CTC and stochastic depth. Following the findings of Section 3, we put intermediate branches at the 6th and 12th layers, and set weight $w=0.66\approx 2/3$. The model is called as the pruning-aware model. After the training, we induce smaller sub-models from the pruning-aware model by removing the layers up to the half, either using an intermediate strategy or iterative strategy in Section 4. We do not fine-tune any sub-models after pruning.

To measure the performance of sub-models, we train individual baseline models with the same depth from scratch. For each baseline, we prepare two configurations: (A) not using intermediate CTC or stochastic depth, and (B) using both of intermediate CTC and stochastic depth following [22]. The baseline B is more challenging for sub-models to reach, as the regularizations are not only useful for pruning but also for improving individual models [22, 27, 28, 5].

All models are trained for 100 epochs for WSJ, and for 50 epochs for AISHELL-1 and TED-LIUM2. We emphasize that the pruning-aware model uses exactly same number of epochs as the individual baseline models do. Thus, it only requires computational cost for training that the 24-layer baseline model does. During testing, we use greedy decoding and do not use any external language models (LMs), enabling fast, parallel and non-autoregressive generation.

Figure 3 shows the word error rates (WERs) for WSJ and TED-LIUM2 and character error rates (CERs) for AISHELL-1. The baselines are displayed as individual marks and pruned sub-models are displayed as lines. We first see the pruning-aware model and 24-layer baseline B have nearly same error rates. This indicates that the proposed pruning method does not harm the full model performance. Also, the half-sized sub-model and 12-layer baseline B also have nearly same error rates. Iterative pruning reaches the individual baseline B models for most of cases. This shows the effectiveness of the suggested method.

We measure the real-time factor (RTF) of the pruning-aware model for each number of layers. We use 1 Xeon CPU of 2.2GHz and 1 P40 GPU respectively. Figure 4 shows that the proposed model is very fast (RTF 0.05 on CPU and 0.005 on GPU), due to lightweight and non-autoregressive properties of CTC, and it can further improve RTF (2.5x on GPU) at the cost of small degradation on WER (18%).

In this section, we show how the change of intermediate CTC proposed at Section 3.2 improves sub-models. We prune 24-layer baseline models of Section 5.1 and compare their performance to the proposed pruning-aware model. Figure 5 shows the WER of sub-models induced by each model, either by intermediate strategy (left) or iterative pruning (right).

We see that baseline A (dashed), without any regularization, does not work well with pruning. It is much improved by baseline B (dotted), but it is still not close to individual baselines. We also see that, for layers between the intermediate and final one, the intermediate sub-model performs poorly.

Finally, the proposed model (solid) gives huge improvement over baseline B, implying that the adjustment of intermediate CTC is very important for pruning, although it does not affect the performance of the full model.

As a preliminary experiment, we also apply in-place distillation [4] to the baseline B. The result (dashdot) shows that, while it improves intermediate sub-models, it does not improve the iteratively pruned sub-models and degrades the full model.