FF2: A Feature Fusion two-stream framework for Punctuation Restoration

For the conventional self-attention[11] mechanism, the computation of scaled dot-product lies:

	$$\displaystyle Attn^{(k)}{(Q,K,V)}=softmax(\frac{J^{(k)}}{\sqrt{d_{emb}/H}})V^{(k)}$$		(1)
	$$\displaystyle J^{(k)}={Q^{(k)}K^{(k)}}^{\mathrm{T}}$$		(2)
	$$\displaystyle Q=x_{i}W_{Q}^{h},K=x_{j}W_{K}^{h},V=x_{j}W_{V}^{h}$$		(3)

Here, $J^{(k)}$ denotes the dot-product between query $Q$ and key $K$ for the $k$-th head, $d_{emb}$ refers to the embedding size, $H$ is the head numbers.

To alleviate the inherent issue (Low-Rank Bottleneck) in multi-head self-attention, we modify the computation of self-attention as a two-stage procedure. First, the attention scores $J^{(k)}$ are obtained via equation 1. Then, these scores are aggregated to update the attention weights via a projection component $\phi$. For the component $\phi$, we introduce a new matrix $P_{\lambda}$ ²²2$P_{\lambda}$ is initialized with standard deviations of $0.1/\sqrt{d_{emb}}$ to share information with sibling heads. Then, we update the $J^{(k)}$ as $\hat{J^{(k)}}$ to represent the new attention weights as below:

$$\displaystyle\begin{pmatrix}\hat{J^{(1)}}\\ \\ \hat{J^{(2)}}\\ \\ \vdots\\ \\ \hat{J^{(h)}}\\ \\ \end{pmatrix}=P_{\lambda}\begin{pmatrix}J^{(1)}\\ \\ J^{(2)}\\ \\ \vdots\\ \\ J^{(h)}\\ \\ \end{pmatrix}+\begin{pmatrix}J^{(1)}\\ \\ J^{(2)}\\ \\ \vdots\\ \\ J^{(h)}\\ \\ \end{pmatrix},P_{\lambda}=\begin{pmatrix}\lambda_{11}&\cdots&\lambda_{1h}\\ \\ \lambda_{21}&\cdots&\lambda_{2h}\\ \\ \vdots&\ddots&\vdots\\ \\ \lambda_{h1}&\cdots&\lambda_{hh}\\ \\ \end{pmatrix}$$

(4)

Here, $\lambda_{hh}$ denotes the learned parameter. We modify the Equation 1 to novel Equation 5 as the interaction scores:

$$Attn^{(k)}{(Q,K,V)}=softmax(\frac{\hat{J^{(k)}}}{\sqrt{d_{emb}/H}})V^{(k)}$$

(5)

Second, we repeat $h$ times the projections in different $Q,K,V$, and concatenate the embeddings. Then, a linear layer is employed to obtain the final matrices $M$:

$$M\left(Q,K,V\right)=Concat\left({\rm Attn}^{(1)},\ldots,{\rm Attn}^{(h)}\right)W+B$$

(6)

Here, $W$ is a learnable parameter and $B$ is the bias. After that, we follow the standard transformer workflow and feed the matrices into the layer normalization and feed-forward layer.

As mentioned, we design a two-stream architecture. One stream (ITE) is based on the pre-trained language model Electra ³³3https://huggingface.co/google/electra-large-discriminator, and another tiny non-pretrained module (TNP) is randomly initialized. The TNP is is significantly smaller than ITE. Here, we replace self-attention with Interaction to encourage message sharing.

During the training, we first obtain two feature vectors which are $C_{i}$ and $C_{m}$ via ITE and TNP modules. Second, $C_{i}$ and $C_{m}$ are concatenated as the input $H\in\mathbb{R}^{n\times(d_{i}+d_{m})}$ of Fusion Layer (FL). FL is just one layer transformer, and the vector $H$ is fed into the FL to obtain the final output sequence $\hat{Y}$. Thus, we use cross-entropy loss $\ell$ as follows:

$$\ell=-\sum_{i=1}^{K}p_{i}\log p_{i}$$

(7)

Here, K is the total number of categories, $p_{i}$ is the predicted probability of label $i$.

The English IWSLT2011[12] is a popular benchmark dataset for punctuation restoration. It contains 2.1M words for training set, 296K words for validation set, and 12626 words for the manual transcription test set. It also contains 12822 words for the actual ASR transcription test set, in which the words are predicted by ASR systems, so it comprises some grammatical errors or wrong words. There are four types of punctuation marks (none, comma, period, and question mark), the distribution of categories in the training dataset is as follows: 85.7% without punctuation, 7.53% with a comma, 6.3% with the period, and 0.47% with question marks. In the following, we will employ precision (P), recall (R), and F1-score (F1) to evaluate FF2 and other approaches.

We compare FF2 with the top-performing approaches. One part is RNN-based and employs a deep recurrent neural network, such as BLSTM-CRF, Teacher-Ensemble, and DRNN-LWMA-pre. The other line is transformer-based, including Self-attention-word-speech, CT-Transformer, and SAPR. SAPR uses a transformer encoder-decoder architecture and views punctuation restoration as a translation task. BERT-base+Adversarial, BERT-large+Transfer, BERT-base+FocalLoss (it utilizes the focal loss not cross-entropy loss to advance the results). RoBERTa-large+augmentation and RoBERTa-large+SCL. RoBERTa-base predicts the tags by multiple context windows, funnel-transformer-xlarge+POS+Fusion+SBS abbreviated as FT+POS incorporates an external POS tagger and fuses its predicted labels into the existing language model to address the problem. ELECTRA-large+Disc-ST named Disc-ST, which introduces a discriminative self-training approach with weighted loss and utilizes the extra dataset to achieve the SOTA performance.[13, 14, 2, 15, 3, 4, 16, 17, 6, 18, 19, 20, 21].

For the two-stream architecture, the Interaction Transformer Encoder is based on the public pre-trained language model (Electra), while tiny non-pretrained (TNP is 6 layers transformer blocks) and Fusion Layer are randomly initialized. Adam optimizer is set to default parameters. The learning rate of Electra and TNP are both 5e-6. Turning to the Fusion Layer, we use 8 attention heads, and the hidden size is 3072. Specifically, the maximum sequence length is 256, the batch size is 8, and the dropout rate in our experiments is set to 0.2. We also introduce R-Drop [22] to act as a regularizer. Moreover, we adopt early stopping with patience of 8 epochs to avoid overfitting. We train and evaluate all the models on the 32GB Tesla V100.

Table 1 presents the results of FF2 compared to top performers. Apparently, the transformer-based approaches are far superior to those RNN-based. Our method accomplishes the best results in the Comma (recall, f1-score), the Question (recall), and the Overall (recall, f1-score), separately. Compared to Disc-ST and FT+POS, FF2 obtains a higher recall and f1-score, and exceeds FT+POS by 1 point in the overall f1-score. Thus, the results validate that our proposed feature fusion two-stream framework is effective.

Some methods advance the performance by incorporating some extra information. For instance, Self-attention-word-speech utilizes both lexical and prosody features. FT+POS employs part-of-speech (POS for short) tools to add additional POS tags as extra knowledge. RoBERTa-large+augmentation construct the data by insertion, substitution, and deletion. Not like them, FF2 leverages the training dataset itself and focuses on the design of architecture. Without extra data, FF2 improves 1.0%, 2.4%, and 12.4% compared to FT+POS, RoBERTa-large+augmentation, and Self-attention-word-speech, respectively.

In comparison with the current SOTA model Disc-ST, Disc-ST advances the performance by utilizing a self-training strategy and extending the training dataset (2M) with a large amount of unlabeled data (30M). Thus, their dataset is nearly 16 times larger than ours. As our model parameters are comparable to Disc-ST’s, the results of FF2 (85.3) compared to Disc-ST (85.2) indicate that our design is robust.

For the results of the ASR transcription test set. Without focal loss or data augmentation, we can find that our approach outperforms the bert-based models by a large margin. Besides, FF2 obtains four best results in total, 82.5% in precision and 83.4 % in f1 of the period, 82.9% in recall of question, and 80.1% in precision of the overall, respectively. Compared to RoBERTa-large+augmentation, FF2 has a slightly lower precision and a competitive F1-score but leads to 3.9 points improvement in recall. One possible reason is that data augmentation for RoBERTa-large+augmentation can benefit the robustness.

We conduct the ablation study and the results are shown in table 1. First, the metric has a large drop (-0.9%) while eliminating the TNP module, which verifies that TNP can capture the feature at hand and benefit ITE module. Second, The metric would decrease 0.3% when removing Interaction, which indicates the modification of self-attention is valid.