Understanding Self-Attention of Self-Supervised Audio Transformers

Adapting the idea of self-supervised learning [1, 2, 3, 4] to continuous speech has received much attention in recent work [5, 6, 7, 8, 9, 10], where Transformer Encoders with multi-head self-attention [11] are pre-trained on a large amount of audio data in a self-supervised scheme. Once pre-trained, they are used to improve various downstream supervised tasks, including phone classification, speaker recognition, SLU, and ASR. Despite the great success of these Self-supervised Audio Transformers (SAT)¹¹1These pre-trained transformer encoders have several different names in their original papers. In this paper we refer to them as SAT for simplicity., their internal attention are often neglected and not explored, as we have little understanding of how they work, or the knowledge they acquire from a large amount of unlabeled data. Understanding how SAT models draw conclusions is crucial for both their improvement and application. In the area of natural language processing (NLP), explaining and interpreting pre-trained black-box models like BERT have been a well-explored topic [12, 13, 14, 15, 16, 17]. However, the analysis of models that are pre-trained on speech has not seen such widespread exploration, and remains an important and challenging endeavor for the speech community.

In this work, we propose to analyze the multi-head self-attention mechanism of SAT through the following methods: visualization, categorization, functionality study, and importance ranking. We found that the self-attentions of SAT models tend to converge into three categories: global attentions, vertical attentions, and diagonal attentions. Diagonal attentions either highly attend to $\pm t$ neighbor or are highly correlated with phoneme boundaries; vertical attentions often concentrate on specific phonemes. As for noisy global attentions, we provide a visualization tool to draw insights about their implicit operations. Through our quantized ranking analysis, we conclude that diagonal attentions outrank the most in terms of importance, followed by vertical attentions. Last but not least, we introduce attention refinement methods which allow us to improve learned representations by moderately removing global attentions or constraining attention span, resulting in a faster inference time and higher performance.

The main ideology of NLP BERT pre-training [1, 2, 3, 4] is to corrupt the input word tokens by randomly masking or permuting them with a probability policy, layers of Transformer Encoder [11] are trained together with a classifier that estimates the masked words at the output. Primarily inspired by this idea, previous works [5, 7, 6, 8, 9, 10] proposed self-supervised learning for audio with Transformer Encoders. In this work, we refer to these types of models as Self-Supervised Audio Transformers, SAT. Unlike BERT where the inputs are discrete text tokens, the inputs of SATs are acoustic features (e.g., MFCC, FBANK, Mel-Spectrogram), which are continuously long and locally smooth. Also, different from BERT where the model is trained by estimating discrete tokens, SATs change the classifier to a feed-forward network and minimize reconstruction error between real frames and predicted frames.

Among all the variants of SATs, we address our focus on SATs that take continuous acoustic frames as input [5, 7, 6], rather than ones that operate on quantized discrete vectors of speech [8, 9, 10]. In our analysis, we particularly consider the framework described in the Mockingjay [5] literature, mainly due to its designs to address the continuously long and locally smooth problem of speech. In Mockinjgay, two techniques of downsampling and consecutive masking are introduced to resolve these issues. Downsampling is applied on input features to adapt SATs to long sequences. To reduce the length of frames by a factor of $R_{factor}$, consecutive frames of $R_{factor}$ amount are reshaped and stacked into one frame [5, 18, 19]. On the other hand, consecutive masking is applied during pre-training to avoid the model from exploiting the local smoothness of acoustic frames. Instead of masking a single frame, consecutive frames of $C_{num}$ are masked to zero. To study the attentions of SAT models, we use the prevailing framework of the LARGE model described in Mockingjay [5], which consists of 12 layers of Transformer Encoders. We train three models on the LibriSpeech [20] train-clean-360 subset with identical settings as in [5], except for $C_{num}\in\{3,6,9\}$, and we name them as M3, M6, M9, respectively.

We start by defining notation for self-attention mechanism and SAT representations. Given a length $T$ sequence of vectors $\bm{x}=x_{1},...,x_{T}\in\mathbb{R}^{d}$, we denote $A^{h}_{u}\in\mathbb{R}^{T\times T}$ as attention weights for all query-key pairs of a head $h$ when propagating an utterance $u$. Hence, $A^{h}_{u}[q,k]\in\mathbb{R}$ is the attention weight of $x_{q}$ attending to $x_{k}$. We use $q$ for timestamp of query; $k$ for timestamp of key, where $1\leq q,k\leq T$. As a result, $A^{h}_{u}[q]\in\mathbb{R}^{T}$ is the attention distribution formed by $x_{q}$, which is a row if we view $A^{h}_{u}$ as a map. When analyzing the representations of a $L$-layer SAT, we denote $\bm{x^{l}}=x^{l}_{1},...,x^{l}_{T}\in\mathbb{R}^{d}$ as the representations of a given layer, where $0\leq l\leq L$ and $\bm{x^{0}}$ represents input features.

We plot out $A^{h}_{u}\in\mathbb{R}^{T\times T}$ as an attention map, where $A^{h}_{u}[0,0]$ starts from the upper-left corner, like Fig 1²²2Due to the page limit, we provide attention maps and supplementary materials in our demo page: https://github.com/leo19941227/Self-Attention-on-SATs . SAT attentions tend to converge into three categories: (1) global: flat attention distributions; (2) vertical: attention maps with vertical lines, and (3) diagonal: attention maps with clear diagonal. Because attention maps of a head are similar across utterances in the sense of three categories, we study self-attention on the basis of head instead of a single attention map. To classify heads into three categories, we define three metrics to quantify a head $h$: globalness $G$, verticality $V$ and diagnality $D$ in equations 1, 2, 3, respectively.

where $\mathbb{H}$ is the standard definition of entropy, and $U$ is a speech corpus. Based on G, V, D, we would have three ranking lists for all heads. If among the three ranking lists, a head has the highest rank based on the list of G, it would be categorized as global, and so on. We use ranking instead of values because the metrics may not have the same numerical scale. Fig 1 shows two attention maps for each category.

Diagonal attentions attend to local neighbors for every query. Some exhibit highly focused like Fig 1(f) and some are block diagonal like Fig 1(e). Interestingly, no SAT contains highly focused diagonal attention at main diagonal. They shift either to the left or right, and larger masking span $C_{num}$ is accompanied by a larger shift, possibly due to SAT models trying to get useful information from further frames. The functionality of block diagonal attentions is discussed in section 5. Vertical attentions like Fig 1(c)(d) always attend to similar locations for all queries given an utterance; global attentions like Fig 1(a)(b) behave randomly. These two categories are discussed in section 6. Finally, we visualize the head distribution^††footnotemark: according to metrics and find the model trained with a larger masking span $C_{num}$ has more global heads. On the contrary, M3 contains the most diagonal heads, suggesting that smaller $C_{num}$ makes SAT focus on local structure more. As for the distribution across layers, all categories of heads are distributed quite uniformly.

There are attention maps of clear block diagonal like Fig 2(a). The borders of blocks might be phoneme boundaries, as illustrated in Fig 2(b). It seems diagonal attention knows phoneme intervals. We conduct phoneme segmentation to examine the correlation. We mainly follow the algorithm proposed in [21], which first calculates a similarity matrix from a sequence of features, containing all pairwise distances between features, and then extract boundary points from the similarity matrix. For segmentation with an attention map, its rows are considered as the feature sequence for computing the similarity matrix^††footnotemark: . Examples of similarity matrices are shown in Fig 2(c)(d) when segmentation features are the attentions map and MFCCs, respectively. We slightly modify the boundary-point-extraction algorithm^††footnotemark: in [21], the modification makes algorithm a little more stable, but only little performance difference is found.

TIMIT [22] is used for evaluating the phoneme segmentation since it provides ground-truth phoneme boundaries. We follow the setup in [23] that uses a small subset of training set as validation to adjust a few algorithm parameters and evaluate on test set. We use a 20ms tolerance window and evaluate with R-value [24] and precision-recall curve. We hand-pick a visually block diagonal head for each of M3, M6, and M9. We choose MFCC as baseline feature since it is the most prevailing feature [21, 25, 26, 27] for segmentation. Little performance difference is found between MFCC and $\bm{x^{0}}$ (Mel-scale spectrogram).

Fig 2(e) verifies the correlation between block diagonal attentions and phoneme boundaries, that attentions clearly surpass MFCC under the same setting. As for R-value, under the strict hit counting scenario [24], MFCC achieves 76.68; M3, M6, M9 achieve 79.99, 78.43, 78.19, respectively. Interestingly, larger masking span $C_{num}$ leads to poorer performance. The reason is that when $C_{num}=3$, masked portions are typically within a phoneme interval, the model learned to utilize features in the same interval to reconstruct. On the other hand, $C_{num}=9$ can sometimes mask an entire phoneme interval, the model then tries to retrieve information beyond the interval.

Worth mentioning, similarity matrices on MFCCs and learned block diagonal attentions have a fundamental difference that the former show high activation on similar but distant frames in Fig 2(d), while the latter are more aware of phoneme neighbor structure. Figures similar to Fig 2(d) are shown^††footnotemark: when we compute similarity matrix on Mel-scale spectrogram or SAT representations, suggesting that despite there are similar frames located far apart, block diagonal heads learned to ignore distant information and focus on neighbor structure.

To study the functionality of global and vertical heads, we propose to align attentions to phoneme relations to see whether some heads focus on looking for specific phoneme relations in the utterances. For a sequence of input features $\bm{x^{0}}$, there exists frame-wise phoneme labels $\bm{y}\in\bm{Y}^{T}$, where $\bm{Y}$ is a predefined phone set. We consider $x^{l}_{q}$ attending to $x^{l}_{k}$ as when observing phoneme $y_{q}$ the head would look for phoneme $y_{k}$. We quantify a phoneme relation $Y_{m}\rightarrow Y_{n}$ inside a head $h$ by summing up all attention weights $A^{h}_{u}[q,k]$ whose phoneme relation $y_{q}\rightarrow y_{k}$ equals $Y_{m}\rightarrow Y_{n}$, over the entire speech corpus. More specifically, we plot a phoneme relation map (PRM) $P_{h}\in\mathbb{R}^{|\bm{Y}|\times|\bm{Y}|}$ by the following equations:

where $1\leq m,n\leq|\bm{Y}|$, $\mathbb{I}$ is indicator function, $P_{h}^{{}^{\prime}},P_{U}\in\mathbb{R}^{|\bm{Y}|\times|\bm{Y}|}$ and $P_{U}$ is the distribution of all possible phoneme relations^††footnotemark: in speech corpus $U$, normalizing the effect of dominating relations like $sil\rightarrow sil$ which appears in all utterances. As a result, positive values in $P_{h}$ represent preference for specific phoneme relations; negative values represent the opposite.

PRMs are plotted using TIMIT [22] with 39 phonemes, and results of several heads are shown^††footnotemark: in Fig 3. Since diagonal heads are interpretable themselves, we focus on vertical and global heads. There are several operations: attending to silence, identity, specific phonemes, and not attending to these (not operations). We observe tendency of vertical heads either focus or neglect specific phonemes for all queries, and we bridge their connections. For later discussion, we use focus and neglect to refer to these types of behaviours.

While a PRM characterizes all phoneme relations of a head, we further define concentration $C_{h}\in\mathbb{R}^{|Y|}$ of a head $h$, where each $C_{h}[n]\in\mathbb{R}$ quantifies the amount of focus (when positive) or neglect (when negative) of a head on specific phoneme $Y_{n}$, over all queries:

Fig 4 verifies the connection between verticality and concentration. We report the maximum focus or neglect for each head. Fig 4(a) points out that heads with high verticality do focus on specific phonemes; Fig 4(b) points out even a slight increase of the verticality V of a head has correlation to concentration, for both focus and neglect. Some low-verticality heads with extreme neglect at the bottom-left of Fig 4(b) are diagonal heads, which always attend to their neighbors dynamically and show extreme neglect for all phonemes.

To evaluate the importance of different attention patterns, we conducted two pruning-based probing experiments. We ablate partial functionality of self-attention directly at inference time in two aspects: (1) ablates an entire head; (2) ablates the visible span for all heads. If an attention pattern is essential, ablating it should exhibit immediate loss in terms of the quality of final representations. We examine representation quality by three probing tasks: spectrogram reconstruction, phoneme classification, and speaker recognition. For the first task, we examine the richness in terms of spectrogram details of refined representations. We reuse the reconstruction head during pre-training and measure L1 loss compared to the original. For the latter two tasks, we examine the usefulness of refined representations on downstream tasks. For phoneme and speaker classifications, we train the downstream models using LibriSpeech [20] train-clean-100 subset and fixed 50k steps. In frame-level setting, we use single-hidden MLP; in utterance-level setting, we use mean-pooling followed by a linear transform. Phoneme classification is conducted under frame-level setting; speaker recognition is under frame-level and utterance-level. Following [5], phoneme labels are obtained by the Montreal Force Aligner [28], and all evaluations are done on the LibriSpeech test-clean subset.

In this paper, we present multiple strategies for analyzing the self-attention mechanism in SATs, including phoneme segmentation, phoneme relation map, and two aspects of pruning. We find several attention functionality and operations. We identify critical attentions and show our visualization tool useful for understanding pruning behavior. Finally, we conclude that we can refine representations and speed up inference time for a given SAT in two aspects: removing global heads or constraining attention span. More materials are provided in our demo page^††footnotemark: .