Don’t Be So Sure! Boosting ASR Decoding via Confidence Relaxation

Recently proposed SSL speech recognition models follows the large body of work showing tremendous success with employing transformer-based encoders for extracting informative features from raw input. In most cases, these models are first pre-trained using vast amounts of unlabeled data and later fine-tuned using labeled data by solving a supervised training objective. The recently proposed Wav2vec 2.0 model (Baevski et al. 2020) is composed of a feature encoder that consists of stacked 1-dimensional convolution layers followed by multiple transformer encoder layers. The feature encoder maps raw audio signals into latent speech representations, which are then fed to the first (lowest) transformer encoder layer. The output of each transformer encoder layer is used as input to the next-in-line transformer encoder layer, utilizing the self-attention mechanism to create contextualized speech representations. The model is pre-trained using a contrastive objective, mapping masked spans of a continuous signal into a discrete set of latent representations using a Gumbel-softmax operation (Jang, Gu, and Poole 2016). The loss is computed by contrasting true quantized latent representations of masked time steps from a set of distractors. Once pre-trained, the model is fine-tuned to speech recognition by projecting each time step’s contextualized representation, created by the top transformer encoder layer, into $C$ classes, where $|C|$ is the target vocabulary size. The fine-tuning uses Connectionist Temporal Classification (CTC) loss (Graves et al. 2006) to optimize the model parameters. The CTC loss enables getting around the alignment mismatch of the speech representations and textual labels, as the input audio length is much longer than the target text sequence and most datasets do not provide an accurate alignment of each audio frame and its corresponding text label. The conditional probability which is used to compute CTC loss is defined as:

$$P_{CTC}(Y\mid X)=\sum_{A\in A_{X,Y}}\prod_{t=1}^{T}{p_{t}(a_{t}\mid X)},$$

(1)

where $A$ is a set of valid alignments for input sequence $X$ and target sequence $Y$, and $p_{t}(a_{t}\mid X)$ is the probability of observing label $a_{t}$ at time $t$.

Different from Wav2vec 2.0 pre-training, Hidden Unit BERT (HuBERT) (Hsu et al. 2021) aims to predict the masked speech representation’s cluster, assigned using K-means applied to the input’s mel-frequency cepstral coefficient features (MFCC). HuBERT uses an offline quantization step to predict, thus, can directly infer target classes rather than relying on a contrastive objective.

Many studies made efforts to shed light on transformer encoder-based models (BERT-like) mechanisms and the information captured within their intermediate layers (Tenney, Das, and Pavlick 2019; Clark et al. 2019), an area of interest dubbed BERTology (Rogers, Kovaleva, and Rumshisky 2020). A key finding is that some features are better represented in certain layers (Shah et al. 2021) and that different layers may encode different input properties. Building on these observations, previous works pointed out the merits of incorporating features encoded in lower layers (Xiong et al. 2018) and the benefits of fusing multiple layers instead of the traditional approach of using the output produced by the single last layer (Kondratyuk and Straka 2019; Su and Cheng 2020). To summarize the information encoded across layers, Ji et al. (2021) propose fusing layer representations by sequentially feeding them to a recurrent layer such as LSTM (Hochreiter and Schmidhuber 1997), and using its final state as a global representation. In previous work, Wang et al. (2020) propose using fusion and pooling methods for Neural Machine Translation. They use an aggregation method similar to ours, however, they aggregate layer representations rather than layer logits, and employ the aggregation on the transformer decoder module.

Different from previous works, our method does not require further training nor additional parameters, and applied on models which have already been pre-trained and fine-tuned. To the best of our knowledge, this is the first work that addresses speech recognition confidence related issues via layer aggregation, and exhibits its effectiveness in both performance and computational costs.

A greedy decoding heuristic would be observing each position as if it is isolated from the rest of the sequence, and selecting the token that maximizes the probability defined as:

$$P_{greedy}(Y\mid X)={\arg\max}_{Y}\prod_{t=1}^{T}{p_{t}(y_{t}\mid X)}.$$

(2)

However, such paradigm might neglect candidates with higher quality, as previous predictions cannot be altered given newly decoded information. For this purpose, a beam search decoder is frequently employed, aiming to maximize the combination of both logits and LM scores:

$$P_{G}(Y\mid X)+\alpha_{1}P_{LM}(Y)+\alpha_{2}|Y|,$$

(3)

where $Y$ is the predicted transcript sequence, $\alpha_{1}$ denotes the weight of the language model probability and $\alpha_{2}$ is the word score $|Y|$ is the length of the sequence.

A network is over-confident when the probability distribution is massed to a single class (Pereyra et al. 2017). In other words, highly confident predictions occurs when the softmax probabilities computed by employing the projection head on H_n have low entropy.

To investigate how confidence evolves throughout the transformer layers, we employ the same projection_head followed by a Softmax operation on the representations emitted by each transformer layer.

While projection_head is trained to project the model’s final representations into the vocabulary space, previous works demonstrated that the same projection layer can be employed on intermediate representations as well (Geva et al. 2021).

Our findings indicate that there is a confidence increase at the top transformer layer, as illustrated in Figure 1. We hypothesize that excessive confidence might prevent beam search from considering promising, yet less confident, predictions and subsequently compromise its performance.

Figure 2 (upper) illustrate the top 4 predicted probabilities at each time step. It is worth noting that the predicted probability distribution at the 5-th position is massed almost entirely at token ”$E$”.

Completing the picture, the predicted token evolution is illustrated in Figure 3. For each transformer layer at each time step, we chose the token with the largest probability, computed using projection_head. This illustrates that a correct prediction may live in intermediate layers rather than the top layer. Thus, utilizing intermediate layers can provide beam search with additional valuable information.

Building on these observations, we propose aggregating the logits of the top $M<N$ transformer layers, forming relaxed and better-informed logits to be used by the beam search decoder. The aggregation of the projected logits can be viewed as a projection of the sum of the $M$ layers. Intuitively, this can be accomplished using the distributive property of the matrices, where $A\times B+A\times C=A\times(B+C)$.

As a consequence of the confidence increase across the layers, a scaling mechanism is required to prevent the top layers from dominating the aggregation product. We apply $L2$ normalization on each layer’s predicted logits, regularizing the impact of overly confident layers.

Formally, the layer aggregated logits can be described as follows:

$$\widehat{logits}(X)=\sum_{n=N-M}^{N}projection\_head(\frac{H_{n}}{||H_{n}||_{2}}).$$

(4)

Finally, the aggregated logits and the top layer’s logits are interpolated, forming the logits used for beam search:

$$\beta\cdot logits(X)+(1-\beta)\cdot\widehat{logits}(X),$$

(5)

where $\beta$ is the aggregation trade-off coefficient.

The following sections provide a comparative study of the aggregation approach effectiveness applied on SSL-ASR models.

We conduct two experimental setups: Baseline and Layer Aggregation. The baseline setup employs the logits produced using the top layer, while the layer aggregation employs the aggregated logits using the corresponding aggregation parameters. Both setups utilize the same decoder, namely, beam search fused with a 4-gram LMs.

We evaluate both setups using the Librispeech ’test-clean’, ’test-other, ’dev-clean’, ’dev-other’ splits, and report the results using the Word Error Rate(WER) and Charater Error Rate(CER) metrics.

Table 2 reports the results of the Baseline and proposed Layer Aggregation methods. The results indicate that the proposed approach either matches or improves WER and CER in all experiments. Layer aggregation is particularly effective in lower-resourced settings, i.e. when the amount of labeled data is 10 hours or less. We find that the optimal number of aggregated layers and aggregation trade-off coefficient depends on the size of the model and the amount of labeled data used for fine-tuning, as demonstrated in Table 1. It can be observed that models fine-tuned on low-resourced labeled data rely more on the aggregation product rather than the logits produced by the top layer.

Figure 5 exhibits the relationship between the number of aggregated layers ($M$) and the weight given to intermediate layers ($1-\beta$) by plotting WER values attained using two Wav2vec 2.0 variants, fine-tuned on different amounts of labeled data. We observed a recurring pattern in the experiments conducted, where models fine-tuned using a small amount of labeled data mostly rely on intermediate layers ($\beta\leq 0.3$), while models fine-tuned using a larger portion of labeled data rely heavily on the top layer’s representations.

We study the performance of the baseline and the layer aggregation approach using increasing beam width levels, as it is well known that computation costs and beam width are tied (Meister, Vieira, and Cotterell 2020). To this end, we experiment with two Wav2vec 2.0 BASE variants, fine-tuned using 100 and 960 hours of Librispeech. We set the aggregation parameters, namely, $\beta$ and $M$ according to those detailed in Table 1.

Figure 4 demonstrates the effect of the beam width on both setups. We observe that while larger beam width typically yields better performance, the aggregation method can match the baseline’s WER levels while using a significantly smaller beam width, hence, reducing inference computational costs. For example, when applied to the ’dev-other’ set, the aggregation method can match the baseline’s performance using a width of 1500, while using a beam width of 400.

To further validate our proposed approach, we compare our method with another relaxation method, namely, Temperature scaling. For this purpose, we added a temperature scaling step (Guo et al. 2017) to the softmax operation of the top layer ($H_{N}$) and calibrated the temperature parameter using Librispeech ’dev-clean’ data. We conducted our experiments using two Wav2vec 2.0 BASE variants, fine-tuned using 100 and 960 hours of Librispeech. Table 3 shows a comparison of the layer aggregation method with standard Temperature scaling. Note, that besides for a single case, the temperature scaling did not show any significant performance gains. Layer aggregation outperformed temperature scaling in all of the examined cases, suggesting that harnessing intermediate layers is beneficial.

Examining the errors of models fine-tuned using different sizes of labeled data, we find that the lower-resourced models’ errors (fined-tuned using 10m, 1h, 10h) are characterized by spelling errors such as dropping characters or transcribing words exactly the way they sound. The errors of resource-rich models mostly occur with rare words. We noted that the aggregation method successfully assisted with both types of errors. Taking rare words errors for example, using the aggregation method ”Phedrus” was corrected to ”Phaedrus”, ”Credius” to ”Critias”, and ”Chelsey” corrected to ”Chelsea”.

Next, we wish to analyze the choices of the number of aggregated layers and the weight given to intermediate layers. The trends shown in Figure 5 show that models fine-tuned using smaller resources place their aggregation weight on intermediate layers regardless of model size, while the weight is shifted towards the top layers as the amount of fine-tuning labeled data increases. These observations are on par with the findings of Pasad et al. (Pasad, Chou, and Livescu 2021) which demonstrated that: (a) semantic and acoustic properties are well encoded at the top layers of models fine-tuned using large amounts of labeled data. (b) low-resource fine-tuning is insufficient for diverging the same top layers from the pre-training objective to speech recognition. The latter gives additional motivation to utilize intermediate layers in low-resource setups and further explains why layer aggregation works.

Furthermore, the proportion of the optimal number of aggregated layers remains the same across model sizes and fine-tuning resources, stretching from 33-41% of the layers in low-resource models to $\sim$50% of the layers in resource-rich models.